"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2024-06-18
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import argparse
import json
import os
import sys
import ipe_param_face_recongnition
from ctypes import *
import struct

sys.path.append("../../../")

from common.python.common import *
from common.python.callback_data_struct import *
from common.python.infer_process import *
from common.python import image_decoder
from common.python import image_encoder


class FaceDetectBox(ctypes.Structure):
    _fields_ = [
        ("xmin", ctypes.c_float),
        ("ymin", ctypes.c_float),
        ("xmax", ctypes.c_float),
        ("ymax", ctypes.c_float),
        ("area", ctypes.c_float),
        ("score", ctypes.c_float),
        ("id", ctypes.c_int),
        ("landmark", ctypes.ARRAY(ctypes.c_float, 10)),
    ]


class FaceDetectInfo(ctypes.Structure):
    _fields_ = [
        ("boxesNum", ctypes.c_uint16),
        ("boxes", ctypes.ARRAY(FaceDetectBox, 256)),
    ]


class bbox(ctypes.Structure):
    _fields_ = [
        ("xmin", ctypes.c_uint32),
        ("ymin", ctypes.c_uint32),
        ("xmax", ctypes.c_uint32),
        ("ymax", ctypes.c_uint32),
        ("score", ctypes.c_float),
        ("id", ctypes.c_int),
        ("label", ctypes.c_char * 64),
    ]


class Box(ctypes.Structure):
    _fields_ = [("boxesnum", ctypes.c_uint32), ("boxes", ctypes.ARRAY(bbox, 256))]


class FaceRect:
    def __init__(self, x, y, w, h):
        self.x = x
        self.y = y
        self.w = w
        self.h = h

    def ensure_even(self, n):
        if n % 2 != 0:
            return n + 1
        return n

    def make_even(self):
        self.x = self.ensure_even(self.x)
        self.y = self.ensure_even(self.y)
        self.w = self.ensure_even(self.w)
        self.h = self.ensure_even(self.h)


def extend_face_box(face_box, image_pixel_w, image_pixel_h):
    w = face_box.xmax - face_box.xmin
    h = face_box.ymax - face_box.ymin

    xmin = max(int(face_box.xmin), 0)
    ymin = max(int(face_box.ymin), 0)
    xmax = min(int(face_box.xmax), int(image_pixel_w - 1))
    ymax = min(int(face_box.ymax), int(image_pixel_h - 1))

    rect = FaceRect(xmin, ymin, xmax - xmin, ymax - ymin)
    rect.make_even()
    return rect


def check_face_box(face_box):
    w = face_box.xmax - face_box.xmin
    h = face_box.ymax - face_box.ymin

    if w <= 20 or h <= 20:
        return False

    return True


def lyn_free_callback(args):
    for i in args:
        error_check(sdk.lyn_free(i))
    return 0


def return_buffers(callback_data) -> int:
    for buffer in callback_data[1]:
        callback_data[0].push(buffer)
    return 0


def main():
    parser = argparse.ArgumentParser(description="Face Detection")

    parser.add_argument(
        "input_image",
        metavar="input_image",
        type=str,
        default="../../../data/retinaface_face_detection_with_image_sample/test.jpeg",
        help="input jpeg file path",
    )
    parser.add_argument(
        "output_image",
        metavar="output_image",
        type=str,
        default="../../../output.jpeg",
        help="output jpeg file path",
    )
    parser.add_argument(
        "device_id", metavar="device_id", type=int, default=0, help="lynxi device id"
    )
    parser.add_argument(
        "--face_detect_model",
        metavar="face_detect_model",
        type=str,
        default="../../../model/retinaface_mb.25_b1_h640_w640_c3_iuint8_ofloat16_0000/Net_0",
        help="face detect model path",
    )
    parser.add_argument(
        "--face_detect_post_plugin",
        metavar="face_detect_post_plugin",
        type=str,
        default="../../../common/plugin/postprocess/cross_compiled_lib/libRetinafacePostPlugin.so",
        help="face detect post process plugin path",
    )
    parser.add_argument(
        "--osd_plugin",
        metavar="osd_plugin",
        type=str,
        default="../../../common/plugin/osd/lib/libOsdPlugin.so",
        help="osd plugin path",
    )

    args = parser.parse_args()

    input_file = args.input_image
    output_file = args.output_image
    device_id = args.device_id
    face_detect_model = args.face_detect_model
    face_detect_post_plugin_path = args.face_detect_post_plugin
    osd_plugin_path = args.osd_plugin

    if not os.path.exists(input_file):
        raise ("error: input jpeg file does not exist.")

    if not is_valid_device(device_id):
        raise ("error: device id is invalid.")

    if not os.path.exists(face_detect_model):
        raise ("error: face detect model does not exist.")

    if not os.path.exists(face_detect_post_plugin_path):
        raise ("error: face detect post process plugin does not exist.")

    if not os.path.exists(osd_plugin_path):
        raise ("error: osd plugin does not exist.")

    # 2. 创建 context 与 stream
    ctx, ret = sdk.lyn_create_context(device_id)
    error_check(ret, "create context")
    ret = sdk.lyn_register_error_handler(default_stream_error_handle)
    error_check(ret, "lyn_register_error_handler")
    stream, ret = sdk.lyn_create_stream()
    error_check(ret, "lyn_create_stream")

    # 3. 初始化解码图片类
    print("image decode")

    image_dec = image_decoder.ImageDecoder()
    image_dec.init(input_file, False)
    imgInfo = image_dec.get_image_dec_info()

    # 4. 同步解码图片到 Device 侧
    decodeImg = sdk.lyn_codec_buf_t()
    decodeImg.size = imgInfo.output.predict_buf_size
    p_decode_buf_out, ret = sdk.lyn_malloc(decodeImg.size)
    error_check(ret)
    decodeImg.data = p_decode_buf_out
    image_dec.decode_image_to_device(stream, decodeImg.data)

    # 5. 加载模型
    faceDetectModelInfo = ModelInfo(face_detect_model)

    # 加载lynxi plugin后处理和osd
    faceDetectPlugin, ret = sdk.lyn_plugin_register(face_detect_post_plugin_path)
    common.error_check(ret, "lyn_plugin_register")
    osdPlugin, ret = sdk.lyn_plugin_register(osd_plugin_path)
    common.error_check(ret, "lyn_plugin_register")

    # 6. 人脸检测IPE处理
    print("detect faces")
    pDetectIpeBuf, ret = sdk.lyn_malloc(faceDetectModelInfo.input_size)
    error_check(ret)
    ipeScrfd = ipe_param_face_recongnition.IpeParamScrfd(
        faceDetectModelInfo.width, faceDetectModelInfo.height
    )
    ipeScrfd.set_img_info(
        imgInfo.output.width,
        imgInfo.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
    )
    ipeScrfd.calc_param(stream, decodeImg.data, pDetectIpeBuf)

    # 7. 调用 APU 推理接口
    pDevDetectApuBuf, ret = sdk.lyn_malloc(faceDetectModelInfo.output_size)
    error_check(ret)
    error_check(
        sdk.lyn_execute_model_async(
            stream,
            faceDetectModelInfo.model,
            pDetectIpeBuf,
            pDevDetectApuBuf,
            faceDetectModelInfo.batch_size,
        )
    )

    # 8. 加载人脸检测post Plugin, 进行模型后处理
    pythonapi.PyCapsule_GetPointer.restype = c_void_p
    pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]

    pDevDetectApuBuf_ptr = pythonapi.PyCapsule_GetPointer(pDevDetectApuBuf, None)

    pDevDetectInfo, ret = sdk.lyn_malloc(ctypes.sizeof(FaceDetectInfo))
    error_check(ret)
    pDevDetectInfo_ptr = pythonapi.PyCapsule_GetPointer(pDevDetectInfo, None)

    post_info = struct.pack(
        "4i2f2P",
        faceDetectModelInfo.height,
        faceDetectModelInfo.width,
        imgInfo.output.height,
        imgInfo.output.width,
        0.25,
        0.45,
        pDevDetectApuBuf_ptr,
        pDevDetectInfo_ptr,
    )

    error_check(
        sdk.lyn_plugin_run_async(
            stream,
            faceDetectPlugin,
            "lynFacePostProcess",
            post_info,
            len(post_info),
        )
    )

    face_detect_size = ctypes.sizeof(FaceDetectInfo)
    pHostDetectInfo_np = np.zeros(face_detect_size, dtype=np.uint8)
    pHostDetectInfo = sdk.lyn_numpy_to_ptr(pHostDetectInfo_np)
    error_check(
        sdk.lyn_memcpy_async(
            stream,
            pHostDetectInfo,
            pDevDetectInfo,
            face_detect_size,
            sdk.lyn_memcpy_dir_t.ServerToClient,
        )
    )
    error_check(sdk.lyn_synchronize_stream(stream))

    pHostDetectInfo_c = pythonapi.PyCapsule_GetPointer(pHostDetectInfo, None)
    c_face_info = ctypes.cast(
        pHostDetectInfo_c, ctypes.POINTER(FaceDetectInfo)
    ).contents

    pHostBoxesInfo = Box()

    if c_face_info.boxesNum > 0:
        pHostBoxesInfo.boxesnum = c_face_info.boxesNum

        for faceIndex in range(c_face_info.boxesNum):
            box = c_face_info.boxes[faceIndex]
            pHostBoxesInfo.boxes[faceIndex].xmin = int(box.xmin)
            pHostBoxesInfo.boxes[faceIndex].ymin = int(box.ymin)
            pHostBoxesInfo.boxes[faceIndex].xmax = int(box.xmax)
            pHostBoxesInfo.boxes[faceIndex].ymax = int(box.ymax)
            pHostBoxesInfo.boxes[faceIndex].score = box.score
            pHostBoxesInfo.boxes[faceIndex].id = box.id

    else:
        # print("no face detected")
        pHostBoxesInfo.boxesnum = 0

    # 9. 加载osd Plugin, 将结果叠加到图像上
    if pHostBoxesInfo.boxesnum > 0:
        pHostBoxesInfo_np = np.frombuffer(
            pHostBoxesInfo,
            dtype=np.uint8,
            count=ctypes.sizeof(Box) // np.dtype(np.uint8).itemsize,
        )
        pHostBoxesInfo_ptr = sdk.lyn_numpy_to_ptr(pHostBoxesInfo_np)
        pDevBoxesInfo, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
        error_check(
            sdk.lyn_memcpy_async(
                stream,
                pDevBoxesInfo,
                pHostBoxesInfo_ptr,
                ctypes.sizeof(Box),
                sdk.lyn_memcpy_dir_t.ClientToServer,
            )
        )

        pDevBoxesInfo_ptr = pythonapi.PyCapsule_GetPointer(pDevBoxesInfo, None)
        decodeImg_data_ptr = pythonapi.PyCapsule_GetPointer(decodeImg.data, None)
        osd_para = struct.pack(
            "1P3i1P4i",
            pDevBoxesInfo_ptr,
            imgInfo.output.width,
            imgInfo.output.height,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
            decodeImg_data_ptr,
            24,
            4,
            2,
            4,
        )
        error_check(
            sdk.lyn_plugin_run_async(
                stream, osdPlugin, "lynDrawBoxAndText", osd_para, len(osd_para)
            )
        )
        error_check(
            sdk.lyn_stream_add_callback(
                stream,
                lyn_free_callback,
                [pDevDetectInfo, pDevBoxesInfo],
            )
        )

    # 9. 进行图片编码
    print("save new image")
    encoder = image_encoder.ImageEncoder()
    encoder.encode_image_and_save(stream, output_file, decodeImg, imgInfo)

    #  10. 销毁资源
    ipeScrfd.destory()
    error_check(sdk.lyn_synchronize_stream(stream))
    error_check(sdk.lyn_destroy_stream(stream))
    error_check(sdk.lyn_plugin_unregister(faceDetectPlugin))
    error_check(sdk.lyn_plugin_unregister(osdPlugin))

    error_check(sdk.lyn_free(decodeImg.data))
    error_check(sdk.lyn_free(pDetectIpeBuf))
    error_check(sdk.lyn_free(pDevDetectApuBuf))

    faceDetectModelInfo.unload_model()

    error_check(sdk.lyn_destroy_context(ctx))
    print("process over")


if __name__ == "__main__":
    main()
